#!/usr/bin/env python3
# coding=utf-8

import rospy
from geometry_msgs.msg import Twist
from geometry_msgs.msg import PoseStamped
from mavros_msgs.srv import *
from mavros_msgs.msg import State
import math
import sys, select, termios, tty
import cap

# 空格：降落
# 5  ：开启offboard模式
# 6  ：解锁，准备起飞
# 7  ：起飞

msg = """
Control Your Turtlebot!
---------------------------
Moving around:
   u    i    o
   j    k    l
   m    ,    .
   a    d   : move left/right
   r    f   : move up/down

q/z : increase/decrease max speeds by 10%
w/x : increase/decrease only linear speed by 10%
e/c : increase/decrease only angular speed by 10%
space key, k : force stop
anything else : stop smoothly
b : switch to OmniMode/CommonMode
CTRL-C to quit
"""
Omni = 0  # 全向移动模式

# 键值对应移动/转向方向
moveBindings = {
    'i': (1, 0),  # 前进
    'o': (1, -1),  # 前右
    'j': (0, 1),  # 左转
    'l': (0, -1),  # 右转
    'u': (1, 1),  # 前左
    ',': (-1, 0),  # 后退
    '.': (-1, 1),  # 后右
    'm': (-1, -1),  # 后左
    'a': (0, -1),  # 左平移
    'd': (0, 1),  # 右平移
}

# 键值对应速度增量
speedBindings = {
    'q': (1.1, 1.1),
    'z': (0.9, 0.9),
    'w': (1.1, 1),
    'x': (0.9, 1),
    'e': (1, 1.1),
    'c': (1, 0.9),
}


def getKey():
    tty.setraw(sys.stdin.fileno())
    rlist, _, _ = select.select([sys.stdin], [], [], 0.1)
    if rlist:
        key = sys.stdin.read(1)
    else:
        key = ''

    termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings)
    return key


speed = 0.2  # 默认移动速度 m/s
turn = 1  # 默认转向速度 rad/s


def vels(speed, turn):
    return "currently:\tspeed %s\tturn %s " % (speed, turn)


sita = 0.0  # 朝向
z = 0
w = 0
zf = 1


def pose_cb(m):
    global sita
    global z
    global w
    global zf
    z = m.pose.orientation.z
    w = m.pose.orientation.w
    if z * w > 0:
        zf = 1
    else:
        zf = -1
    sita = 2 * math.acos(w) * 180 / math.pi



current_state = State()


def state_cb(state):
    global current_state
    current_state = state


if __name__ == "__main__":
    settings = termios.tcgetattr(sys.stdin)

    rospy.init_node('turtlebot_teleop')
    pub = rospy.Publisher('mavros/setpoint_velocity/cmd_vel_unstamped', Twist, queue_size=5)
    rospy.Subscriber('mavros/local_position/pose', PoseStamped, pose_cb)
    rospy.Subscriber('mavros/state', State, state_cb)

    setModeServer = rospy.ServiceProxy('mavros/set_mode', SetMode)
    armServer = rospy.ServiceProxy('/mavros/cmd/arming', CommandBool)

    x = 0  # 前进后退方向
    y = 0  # 左右移动方向
    z = 0  # 上下移动方向
    th = 0  # 转向/横向移动方向
    count = 0  # 键值不再范围计数
    target_speed = 0  # 前进后退目标速度
    target_z_speed = 0  # 上下运动目标速度
    target_turn = 0  # 转向目标速度
    target_side_speed = 0  # 左右平移目标速度
    control_speed = 0  # 前进后退实际控制速度
    control_z_speed = 0  # 上下运动实际控制速度
    control_turn = 0  # 转向实际控制速度
    control_side_speed = 0  # 左右平移实际控制速度
    try:
        print(msg)
        print(vels(speed, turn))
        while (1):
            key = getKey()

            if key == 'i':  # 前进
                count = 0
                x = 1
                y = 0
                z = 0
            elif key == ',':  # 后退
                count = 0
                x = -1
                y = 0
                z = 0
            elif key == 'j':  # 往左转
                count = 0
                th = 1
                y = 0
                z = 0
            elif key == 'l':  # 往右转
                count = 0
                th = -1
                y = 0
                z = 0
            elif key == 'r':  # 上升
                count = 0
                z = 2
            elif key == 'f':  # 下降
                count = 0
                z = -2
            elif key == 'a':  # 左平移
                count = 0
                y = -1
            elif key == 'd':  # 右平移
                count = 0
                y = 1
            elif key in speedBindings.keys():
                speed = speed * speedBindings[key][0]
                turn = turn * speedBindings[key][1]
                count = 0
                print(vels(speed, turn))

            elif key == 'k':
                x = 0
                y = 0
                z = 0
                th = 0
                control_speed = 0
                control_z_speed = 0
                control_turn = 0
                control_side_speed = 0

            elif key == ' ':
                print("Vehicle Land")
                setModeServer(custom_mode='AUTO.LAND')
            elif key == '5':
                if current_state.mode != "OFFBOARD":
                    setModeServer(custom_mode='OFFBOARD')
                    print("Offboard enabled")
            elif key == '6':
                armServer(True)
                print("Vehicle armed")
            elif key == '7':
                print("Vehicle Takeoff")
                setModeServer(custom_mode='AUTO.TAKEOFF')
            else:
                count = count + 1
                if count > 4:
                    x = 0
                    y = 0
                    z = 0
                    th = 0
                if (key == '\x03'):
                    break

            target_speed = speed * x
            target_z_speed = speed * z
            target_turn = turn * th
            target_side_speed = speed * y

            if target_speed > control_speed:
                control_speed = min(target_speed, control_speed + 0.1)
            elif target_speed < control_speed:
                control_speed = max(target_speed, control_speed - 0.1)
            else:
                control_speed = target_speed

            if target_z_speed > control_z_speed:
                control_z_speed = min(target_z_speed, control_z_speed + 0.1)
            elif target_z_speed < control_z_speed:
                control_z_speed = max(target_z_speed, control_z_speed - 0.1)
            else:
                control_z_speed = target_z_speed

            if target_turn > control_turn:
                control_turn = min(target_turn, control_turn + 0.5)
            elif target_turn < control_turn:
                control_turn = max(target_turn, control_turn - 0.5)
            else:
                control_turn = target_turn

            if target_side_speed > control_side_speed:
                control_side_speed = min(target_side_speed, control_side_speed + 0.1)
            elif target_side_speed < control_side_speed:
                control_side_speed = max(target_side_speed, control_side_speed - 0.1)
            else:
                control_side_speed = target_side_speed

            y_sita = math.sin(sita / 180 * math.pi)
            if y_sita < 0:
                y_sita = -y_sita
            y_sita = y_sita * zf

            twist = Twist()
            twist.linear.x = control_speed * math.cos(sita / 180 * math.pi)
            twist.linear.y = control_speed * y_sita + control_side_speed  # 包含左右平移
            twist.linear.z = control_z_speed
            twist.angular.x = 0
            twist.angular.y = 0
            twist.angular.z = control_turn

            pub.publish(twist)

    except Exception as e:
        print(e)

    finally:
        twist = Twist()
        twist.linear.x = 0
        twist.linear.y = 0
        twist.linear.z = 0
        twist.angular.x = 0
        twist.angular.y = 0
        twist.angular.z = control_turn
        pub.publish(twist)

    termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings)
